from pprint import pprint

from flask import Flask, request

app = Flask(__name__)

def fun_add(a: int, b: int) -> int:
    return a + b

def fun_mul(a: int, b: int) -> int:
    return a * b

def fun_square(a: int) -> int:
    return a**2

def fun_greet() -> str:
    return 'Hello World!'

def fun_raise(a: int) -> int:
    raise RuntimeError({"code":a, "msg": "I raise!!!"})

def fun_critical() -> None:
    raise Exception("I am critical!")

methods_avail = {
    "add": fun_add,
    "mul": fun_mul,
    "square": fun_square,
    "greet": fun_greet,
    "raise": fun_raise,
    "critical": fun_critical
}

@app.route("/jsonrpc", methods=['GET', 'POST'])
def json_rpc():
    print(request.get_data())

    # step 1. validate request
    if not request.is_json:  # validate Content-Type
        return {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}, "id": None}
    content = request.get_json(silent=True)
    if content is None:  # parse fail or empty json
        return {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}, "id": None}

    if isinstance(content, list):
        return [handle_rpc(call) for call in content]
    elif isinstance(content, dict):
        return handle_rpc(content)


def handle_rpc(content: dict) -> dict:
    pprint(content)

    # step 2. validate json structure
    json_has_version = "jsonrpc" in content and content["jsonrpc"] == "2.0"
    json_has_method = "method" in content and content["method"] is not None
    json_has_params = "params" in content and content["params"] is not None
    json_has_id = "id" in content and content["id"] is not None

    json_id = content["id"] if json_has_id else None

    if not(json_has_version and json_has_method) :
        return {"jsonrpc": "2.0", "error": {"code": -32600, "message": "Invalid Request"}, "id": json_id}


    # step 3. find called function
    if not content["method"] in methods_avail.keys():
        return {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found"}, "id": json_id}
    method = methods_avail[content["method"]]


    # step 4. check request param for function
    result = None
    try:
        if not json_has_params:
            result = method()
        elif isinstance(content["params"], dict):
            result = method(**content["params"])  # kwargs
        elif isinstance(content["params"], list):
            result = method(*content["params"])  # positional
        elif content["params"] is not None:
            result = method(content["params"])  # single param
        else:
            raise TypeError


    # step 5. error handling
    except TypeError as e:
        return {"jsonrpc": "2.0", "error": {"code": -32602, "message": "Invalid Params"}, "id": json_id}
    except RuntimeError as e:  # raise this in your function
        if e.args[0] is not None and isinstance(e.args[0], dict):
            err_info = e.args[0]
            err_code = -32000 - err_info["code"] if "code" in err_info else -32000
            err_msg = err_info["msg"] if "msg" in err_info else "shit happened"
            if "data" in err_info:
                return {"jsonrpc": "2.0", "error": {"code": err_code, "message": err_msg, "data": err_info["data"]}, "id": json_id}
            return {"jsonrpc": "2.0", "error": {"code": err_code, "message": err_msg}, "id": json_id}
        else:
            return {"jsonrpc": "2.0", "error": {"code": -32000, "message": "Server error"}, "id": json_id}
    except Exception as e:
        return {"jsonrpc": "2.0", "error": {"code": -32603, "message": "Internal error"}, "id": json_id}


    # finally return
    return {
        "jsonrpc": 2.0,
        "result": result,
        "id": json_id,
    }

@app.route('/')
def hello_world():  # put application's code here
    return 'Hello World!'


if __name__ == '__main__':
    app.run()
